"""
========================================================================
FILE: mapper.py

LAST MODIFIED: 01 September 2009

DESCRIPTION: 
class to create maps mapping ensemble parameters to element point 
parameters and vice-versa.
========================================================================
"""
"""
BEGIN LICENSE BLOCK 
Version: MPL 1.1/GPL 2.0/LGPL 2.1

The contents of this file are subject to the Mozilla Public License Version
1.1 (the "License"); you may not use this file except in compliance with
the License. You may obtain a copy of the License at
http://www.mozilla.org/MPL/

Software distributed under the License is distributed on an "AS IS" basis,
WITHOUT WARRANTY OF ANY KIND, either express or implied. See the License
for the specific language governing rights and limitations under the
License.

The Original Code is fieldwork.

The Initial Developer of the Original Code is Ju Zhang.
Portions created by the Initial Developer are Copyright (C) 2009
the Initial Developer. All Rights Reserved.

Contributor(s):

Alternatively, the contents of this file may be used under the terms of
either the GNU General Public License Version 2 or later (the "GPL"), or
the GNU Lesser General Public License Version 2.1 or later (the "LGPL"),
in which case the provisions of the GPL or the LGPL are applicable instead
of those above. If you wish to allow use of your version of this file only
under the terms of either the GPL or the LGPL, and not to allow others to
use your version of this file under the terms of the MPL, indicate your
decision by deleting the provisions above and replace them with the notice
and other provisions required by the GPL or the LGPL. If you do not delete
the provisions above, a recipient may use your version of this file under
the terms of any one of the MPL, the GPL or the LGPL.

END LICENSE BLOCK
"""

import scipy

class mapper:
	""" Class to create maps mapping ensemble parameters to element point 
	parameters and vice-versa.
	"""
	def __init__( self, debug = 0 ):
		self.mesh = None
		self._ensemble_to_element_map = {}	# { global_point_number: { element_number: {point_number: weighting}, ... } }
		self._element_to_ensemble_map = {}	# { element_number: { point_number: [[ global_point_number, [weightings] ], [...], ...] } }
		self.number_of_ensemble_points = 0
		self.number_of_dof = 0
		self._custom_ensemble_order = None	# { self generated number: user assigned number }
		self.debug = debug
		
		return
	
	#==================================================================#	
	def set_parent_field( self, parent_field ):
		""" set new parent field on which mapping is based.
		parent field should contain toplogical, and basis information 
		about itself and child fields.
		"""
		
		self.field = parent_field
		self.number_of_ensemble_points = self.field.mesh.get_number_of_ensemble_points()
		
		if not self.number_of_ensemble_points:
			print 'ERROR: mapper.set_parent_field: empty parent mesh'
			return
		else:
			# initialise ensemble to element map
			for i in range( self.number_of_ensemble_points ):
				self._ensemble_to_element_map[i] = {}
			
			# initialise element to ensemble map
			for i in range( len( self.field.mesh.elements ) ):
				# initialise element entries
				self._element_to_ensemble_map[i] = {}
				for j in range( self.field.mesh.elements[i].get_number_of_ensemble_points() ):
					# initialise element point entries for each element
					self._element_to_ensemble_map[i][j] = []
					
			return 1
	
	#==================================================================#	
	def do_mapping( self ):
		""" map the current mesh topology
		"""
		
		gp = 0	# global points number
		assigned = []	# list of element points already assigned a global number
		ep = self.field.mesh.connectivity.keys()	# list of all element points and hanging points
		ep.sort()
		i = 0
		
		if self.debug:
			print 'sorted element points:', ep
		
		# account for points connected to hanging nodes
		while ep[i][0] == -1:
			assigned += self.field.mesh.connectivity[ep[i]]
			i += 1
		
		# map conventional points
		for j in range(i, len( ep ) ):
			
			if self.debug:
				print 'mapping ep:', ep[j]
			
			# loop through each element point in the connectivity dict
			(e, p) = ep[j]
			# check if point has already been assigned a global number
			if (e, p) not in assigned:
				
				if self.debug:
					print 'assigning global', gp, 'to', [(e,p)] + self.field.mesh.connectivity[(e,p)]
				
				# update ensemble to element map
				try:
					self._ensemble_to_element_map[gp][e][p] = 1.0
				except KeyError:
					self._ensemble_to_element_map[gp][e] = { p: 1.0 }

				# update element to ensemble map
				self._element_to_ensemble_map[e][p].append( [ gp, [1.0] ] )
				assigned.append( (e,p) )
				
				# map points connected to the current element point
				for [ec, pc] in self.field.mesh.connectivity[(e,p)]:
					# update ensemble to element map
					try:
						self._ensemble_to_element_map[gp][ec][pc] = 1.0
					except KeyError:
						self._ensemble_to_element_map[gp][ec] = { pc: 1.0 }
					
					# update element to ensemble map
					self._element_to_ensemble_map[ec][pc].append( [ gp, [1.0] ] )
					assigned.append( (ec, pc) )
					
				if self.debug:
					print 'assigned:', assigned
				gp += 1
		
		if self.debug:
			print 'element to ensemble map:', self._element_to_ensemble_map
		
		# map hanging points
		if self.debug:
			print 'i:', i
		# for each hanging point
		for h in range( i ):
			hp = self.field.mesh.hanging_points[h]
			
			# get its host_element and element_coord
			host_element = hp.get_host_element()
			host_element_c = hp.get_element_coordinates()
			
			# find the ensemble points associated with the host element
			host_gp = []
			for element_point in self._element_to_ensemble_map[host_element].keys():
				for [g, w] in self._element_to_ensemble_map[host_element][element_point]:
					host_gp.append(g)
				
			# calculate weightings using basis
			# if local
			weights = self.field.basis.eval( host_element_c )
			if self.debug:
				print 'host_element_c:', host_element_c, 'weights:', weights
				print 'host_gp:', host_gp
			
			if isinstance( weights, int ):
				print 'ERROR: mapper._map_ensemble_to_element: unable to evaluate weights'
				return
			else:
				# update maps for connected element points
				for [e,p] in self.field.mesh.connectivity[ ep[h] ]:
					for host_gp_i in range( len( host_gp ) ):
						# update ensemble to element map
						try:
							self._ensemble_to_element_map[ host_gp[host_gp_i] ][e][p] = weights[ host_gp_i ]
						except KeyError:
							self._ensemble_to_element_map[ host_gp[host_gp_i] ][e] = { p: weights[ host_gp_i ] }

						# update element to ensemble map
						self._element_to_ensemble_map[e][p].append( [ host_gp[host_gp_i], [weights[ host_gp_i ]] ] )
			
		return 1
	
	#==================================================================#
	def remove_element( self, element_number ):
		""" removes element from the _element_to_ensemble_map
		"""
		
		if element_number in self._element_to_ensemble_map.keys():
			
			# get element point to ensemble points mapping
			ep_map = self._element_to_ensemble_map[ element_number ]
			ensemble_points = []
			for element_point in ep_map.values():
				ensemble_points += [ p[0] for p in element_point ]
			
			# remove entries in ensemble_to_element_map	
			for ensemble_point in ensemble_points:
				del self._ensemble_to_element_map[ensemble_point][element_number]
			
			# remove mapping in _element_to_ensemble_map
			del self._element_to_ensemble_map[ element_number ]
			
			return 1
		else:
			print 'ERROR: mapper.remove_element: element', element_number, 'does not exist'
			return
	
	#==================================================================#
	def get_ensemble_point_element_points( self, ensemble_point_number ):
		""" returns a list of ( element n, element point n ) mapped to the
		given ensemble_point_number
		"""
		
		try:
			element_map = self._ensemble_to_element_map[ensemble_point_number]
		except KeyError:
			print 'ERROR: mapper.get_ensemble_point_number: invalid ensemble point number'
			return None
		else:
			element_points = []
			for element in element_map:
				for point in element_map[element]:
					element_points.append( (element, point) )
					
			return element_points
			
	
	#==================================================================#
	def get_element_parameters( self, element_number, parameters ):
	 	""" Uses mapping to return the element parameters for the
		required element. 
		"""
		
		# get mapping for specified element
		try:
			element_map = self._element_to_ensemble_map[ element_number ]
		except IndexError:
			print 'ERROR: mapper.get_element_parameters: invalid element_number'
			return None
		
		if len( parameters ) != self.get_number_of_ensemble_points():
			print 'ERROR: mapper.get_element_parameters: incorrect number of parameters. Require', self.get_number_of_ensemble_points()
			return None
		
		# get list of element points
		points = element_map.keys()
		points.sort()
		
		element_parameters = []
		# for each element point
		for element_point in points:
			mapped_values = []
			# for each global point mapped to this element point
			for [g, w] in element_map[ element_point ]:
				
				# if custom ensemble point ordering, map g to the right ensemble point
				if self._custom_ensemble_order:
					g = self._custom_ensemble_order[g]
				
				# calculated the weighted contribution of this global point
				mapped_values.append( scipy.multiply( parameters[g], w ) )
			
			# sum contributions and append element_point parameter to element_parameters
			element_parameters.append( list( scipy.sum( mapped_values, axis = 0 ) ) )
		
		return element_parameters
	
	#==================================================================#	
	def get_number_of_ensemble_points( self ):
		return self.number_of_ensemble_points
	
	#==================================================================#		
	def get_number_of_dof( self ):
		return self.number_of_dof
	
	#==================================================================#		
	def set_custom_ensemble_ordering( self, ordering ):
		""" defines custom numbering order to ensemble points
		"""
		
		# check for correct length
		if len( ordering ) != self.number_of_ensemble_points:
			print 'ERROR: mapper.set_custom_ensemble_ordering: wrong number of entries in new ordering. Require:', self.number_of_ensemble_points, ', got:', len( ordering )
			return
		else:
			self._custom_ensemble_order = dict(ordering)
			return 1
	
